package middleware

import (
	"strings"

	"goadmin-huasen-portal/serializer"
	"goadmin-huasen-portal/util"

	cmap "github.com/endymx/concurrent-map"
	"github.com/gin-gonic/gin"
	"github.com/golang-jwt/jwt"
	"golang.org/x/time/rate"
)

var userLimit = cmap.NewWithCustomShardingFunction[int64, *rate.Limiter](func(key int64) uint32 {
	return uint32(key) ^ uint32(key>>32)
})

type UserInfoMiddleware struct {
	AccessSecret string
}

func NewUserInfoMiddleware(accessSecret string) *UserInfoMiddleware {
	return &UserInfoMiddleware{
		AccessSecret: accessSecret,
	}
}

func (m *UserInfoMiddleware) Handle() gin.HandlerFunc {
	return func(c *gin.Context) {

		parts := strings.SplitN(c.GetHeader("Authorization"), " ", 2)
		if len(parts) != 2 || parts[0] != "Bearer" {
			c.JSON(400, serializer.Err(serializer.CodeCheckLogin, "登录过期,请重新登录", nil))
			c.Abort()
			return
		}
		t := parts[1]
		tok, err := jwt.Parse(t, func(token *jwt.Token) (interface{}, error) {
			return []byte(m.AccessSecret), nil
		})
		if err != nil || !tok.Valid {
			util.Log().Error("Invalid token :", err.Error())
			c.JSON(400, serializer.Err(serializer.CodeCheckLogin, "登录过期,请重新登录", nil))
			c.Abort()
			return
		}

		claims, ok := tok.Claims.(jwt.MapClaims)
		if len(claims) == 0 || !ok {
			c.JSON(400, serializer.Err(serializer.CodeCheckLogin, "登录过期,请重新登录", nil))
			c.Abort()
			return
		}
		uid := int64(claims["uid"].(float64))
		if uid <= 0 {
			c.JSON(400, serializer.Err(serializer.CodeCheckLogin, "登录过期,请重新登录", nil))
			c.Abort()
			return
		}
		if _, ok := userLimit.Get(uid); !ok {
			userLimit.Set(uid, rate.NewLimiter(5, 10))
		}
		if limit, _ := userLimit.Get(uid); !limit.Allow() {
			c.JSON(400, serializer.Err(serializer.CodeServerErr, "访问过快, 请稍后再试", nil))
			c.Abort()
			return
		}

		c.Set("uid", uid)
		c.Next()
	}
}
